package com.example.day04.controller.union;

import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ConnectDemo2 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1); // 设置并行度为 1

        // 创建第一条数据流 source1，包含 Tuple2<Integer, String>
        DataStreamSource<Tuple2<Integer, String>> source1 = env.fromElements(
                Tuple2.of(1, "a1"),
                Tuple2.of(2, "b"),
                Tuple2.of(3, "c")
        );

        // 创建第二条数据流 source2，包含 Tuple3<Integer, String, Integer>
        DataStreamSource<Tuple3<Integer, String, Integer>> source2 = env.fromElements(
                Tuple3.of(1, "aa1", 1),
                Tuple3.of(1, "aa2", 2),
                Tuple3.of(2, "bb", 1),
                Tuple3.of(3, "cc", 1)
        );

        // 使用 connect() 方法连接两条数据流
        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> connect = source1.connect(source2);

        ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>> tuple2Tuple3ConnectedStreams = connect.keyBy(s1 -> s1.f0, s2 -> s2.f0);
        SingleOutputStreamOperator<String> process = tuple2Tuple3ConnectedStreams.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>() {

            // 每个任务，每个分区内存一个哈希映射，用于缓存数据
            Map<Integer, List<Tuple2<Integer, String>>> s1Cache = new HashMap<>();
            Map<Integer, List<Tuple3<Integer, String, Integer>>> s2Cache = new HashMap<>();

            @Override
            public void processElement1(Tuple2<Integer, String> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {

                Integer id = value.f0;
                if (!s1Cache.containsKey(id)) { // 如果 key 不存在
                    // 1.1 如果 key 不存在，说明是该 key 的第一条数据，初始化，放入 map
                    List<Tuple2<Integer, String>> s1Values = new ArrayList<>(); // 初始化一个空列表
                    s1Values.add(value); // 添加当前值
                    s1Cache.put(id, s1Values); // 把这个列表放到哈希映射里
                } else {
                    // 1.2 key 存在，不是该 key 的第一条数据，直接添加到 value 的 list 中
                    s1Cache.get(id).add(value); // 获取已存在的列表并添加当前值
                }

                // TODO 2. 去 s2Cache 中查找是否有 id 匹配上的，匹配上就输出，没有就不输出
                if (s2Cache.containsKey(id)) {
                    for (Tuple3<Integer, String, Integer> s2Element : s2Cache.get(id)) {
                        out.collect("s1:" + value + "=======" + "s2:" + s2Element);
                    }
                }
            }


            @Override
            public void processElement2(Tuple3<Integer, String, Integer> value, CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, Integer>, String>.Context ctx, Collector<String> out) throws Exception {
                Integer id = value.f0;
                if (!s2Cache.containsKey(id)) { // 如果 key 不存在
                    // 1.1 如果 key 不存在，说明是该 key 的第一条数据，初始化，放入 map
                    List<Tuple3<Integer, String, Integer>> s2Values = new ArrayList<>(); // 初始化一个空列表
                    s2Values.add(value); // 添加当前值
                    s2Cache.put(id, s2Values); // 把这个列表放到哈希映射里
                } else {
                    // 1.2 key 存在，不是该 key 的第一条数据，直接添加到 value 的 list 中
                    s2Cache.get(id).add(value); // 获取已存在的列表并添加当前值
                }

                // TODO 2. 去 s2Cache 中查找是否有 id 匹配上的，匹配上就输出，没有就不输出
                if (s1Cache.containsKey(id)) {
                    for (Tuple2<Integer, String> s1Element : s1Cache.get(id)) {
                        out.collect("s1:" + s1Element + "=======" + "s2:" + value);
                    }
                }
            }
        });
        process.print("zhao");
        env.execute();
    }
}
